import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 网格分辨率参数
parm = [8, 16, 32, 64]

# 创建一个大图，包含多个子图
fig = plt.figure(figsize=(15, 10))

# 遍历每个网格分辨率
for idx, resolution in enumerate(parm):
    # 读取数据
    filename = f'Problem2_parm{resolution}.txt'
    data = np.loadtxt(filename)

    # 提取x, y, z坐标
    x = data[:, 0]
    y = data[:, 1]
    z = data[:, 2]

    # 创建网格
    # 由于部分点被移除，需要重新构造网格
    m = resolution + 2  # 网格分辨率（包括边界点）
    x_grid = np.linspace(0, 1, m)
    y_grid = np.linspace(0, 1, m)
    x, y = np.meshgrid(x_grid, y_grid)

    # 将z值填充到网格中
    z_grid = np.full((m, m), np.nan)  # 初始化网格为NaN
    for i in range(len(data)):
        xi = int(np.round(data[i, 0] * (m - 1)))  # 计算x索引
        yi = int(np.round(data[i, 1] * (m - 1)))  # 计算y索引
        z_grid[yi, xi] = data[i, 2]  # 填充z值

    # 添加子图
    ax = fig.add_subplot(2, 2, idx + 1, projection='3d')  # 2行2列的子图布局

    # 绘制表面图
    surf = ax.plot_surface(x, y, z_grid, cmap='viridis', vmin=1, vmax=6)

    fig.colorbar(surf, ax=ax, shrink=0.6, aspect=10)
    # 设置标签和标题
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('U')
    ax.set_title(f'Grid Resolution: {resolution}x{resolution}')

# 调整布局
plt.tight_layout()
plt.savefig('../../figure/F1_Disk_N.png')
# 显示图形
plt.show()